# 改完了
import mindspore.nn as nn
import mindspore
import mindspore.ops as ops
from .module import *
# from module import *
class DepthHead(nn.Cell):
    def __init__(self, input_dim=256, hidden_dim=128, scale=False):
        super(DepthHead, self).__init__()
        self.scale = scale
        self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1, pad_mode='pad')
        self.conv2 = nn.Conv2d(hidden_dim, 1, 3, padding=1, pad_mode='pad')
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout2d(p=0.1)


    # def forward(self, x_d, act_fn=torch.tanh):
    def construct(self, x_d, act_fn=ops.tanh):
        out = self.conv2(self.relu(self.conv1(x_d)))
        if self.training and self.dropout is not None:
            out = self.dropout(out)
        # out = nn.ReLU(out)
        # out = torch.clip(out, max = 2)
        # return out - 1
        return act_fn(out)

class ConvGRU(nn.Cell):
    def __init__(self, hidden_dim=128, input_dim=192+128):
        super(ConvGRU, self).__init__()
        self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1, pad_mode='pad')
        self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1, pad_mode='pad')
        self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1, pad_mode='pad')

    def construct(self, h, *x_list):
        x =  ops.cat(x_list, axis=1)
        hx = ops.cat([h, x], axis=1)
        z =  ops.sigmoid(self.convz(hx))
        r =  ops.sigmoid(self.convr(hx))
        q =  ops.tanh(self.convq(ops.cat([r*h, x], axis=1)))
        h = (1-z) * h + z * q
        return h

class UpMaskNet(nn.Cell):
    def __init__(self, hidden_dim=128, ratio=8):
        super(UpMaskNet, self).__init__()
        self.mask = nn.SequentialCell( 
            nn.Conv2d(hidden_dim, hidden_dim*2, 3, padding=1, pad_mode='pad'),
            nn.ReLU(),
            nn.Conv2d(hidden_dim*2, ratio*ratio*9, 1, padding=0, pad_mode='pad'))

    def construct(self, feat):
        # scale mask to balence gradients
        mask = .25 * self.mask(feat)
        return mask

class ProjectionInput(nn.Cell):
    def __init__(self, cost_dim, hidden_dim, context_dim, out_chs, depth_num=1, G=8):
        super().__init__()

        self.convc1 = nn.Conv2d(cost_dim, hidden_dim, 1, padding=0, pad_mode='pad')
        self.convc2 = nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1, pad_mode='pad')
        self.convd1 = nn.Conv2d(depth_num, hidden_dim, 7, padding=3, pad_mode='pad')
        self.convd2 = nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1, pad_mode='pad')
        self.convd = nn.Conv2d(hidden_dim + hidden_dim, hidden_dim-context_dim, 3, padding=1, pad_mode='pad')
        # self.convc = nn.Conv2d(hidden_dim, hidden_dim//2, 1, padding=0)
        # self.out_chs = hidden_dim//2
        self.convc = nn.Conv2d(hidden_dim, hidden_dim, 1, padding=0, pad_mode='pad')
        self.out_chs = hidden_dim
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout2d(p=0.1)

    def construct(self, disp, cost, context):
        cor = self.relu(self.convc1(cost))
        cor = self.relu(self.convc2(cor))
        dfm = self.relu(self.convd1(disp))
        dfm = self.relu(self.convd2(dfm))
        cor_dfm = ops.cat([cor, dfm], axis=1)
        cor_dfm = self.convd(cor_dfm)
        cor_dfm = self.convc(ops.cat([cor_dfm, context], axis=1))
        out_d = self.relu(cor_dfm)

        if self.training and self.dropout is not None:
            out_d = self.dropout(out_d)
        return out_d

class BasicUpdateBlock(nn.Cell):
    def __init__(self, hidden_dim=128, cost_dim=256, ratio=8, context_dim=64 ,UpMask=False, Inverse=False, cost_num=1, G=8):
        super(BasicUpdateBlock, self).__init__()
        self.encoder = ProjectionInput(cost_dim=cost_dim*cost_num, hidden_dim=hidden_dim, context_dim=context_dim, out_chs=hidden_dim, G=G)
        self.depth_gru = ConvGRU(hidden_dim=hidden_dim, input_dim=self.encoder.out_chs)
        self.depth_head = DepthHead(hidden_dim, hidden_dim=hidden_dim, scale=False)
        self.UpMask = UpMask
        self.Inverse = Inverse
        self.mask = nn.SequentialCell(
            nn.Conv2d(hidden_dim, hidden_dim*2, 3, padding=1, pad_mode='pad'),
            nn.ReLU(),
            nn.Conv2d(hidden_dim*2, ratio*ratio*9, 1, padding=0, pad_mode='pad'))

    def construct(self, net, depth_cost_func, inv_depth, context, seq_len=4, scale_inv_depth=None):
        inv_depth_list = []
        mask_list = []
        print("enter BasicUpdateBlock")

        for i in range(seq_len):
            # TODO detach()
            inv_depth = ops.stop_gradient(inv_depth)
            print("enter encoder")
            input_features = self.encoder(inv_depth, depth_cost_func(scale_inv_depth(inv_depth)[1], iter=i), context)
            print("enter depth_gru")
            # inp_i = ops.cat([context, input_features], dim=1)
            net = self.depth_gru(net, input_features)
            print("enter depth_head")
            delta_inv_depth = self.depth_head(net)
            inv_depth_test = inv_depth + delta_inv_depth
            # if inv_depth_test.mean() != inv_depth_test.mean():
            #     print("error2")
            #     print(delta_inv_depth.mean())
            #     print(net.mean())
            #     print(inv_depth.mean())
            #     print(input_features.mean())
            #     print(depth_cost_func(scale_inv_depth(inv_depth)[1]).mean())
            inv_depth = inv_depth_test
            inv_depth_list.append(inv_depth)
            if self.UpMask and i == seq_len - 1 :
                mask = .25 * self.mask(net)
                mask_list.append(mask)
            else:
                mask_list.append(inv_depth)
        return net, mask_list, inv_depth_list

if __name__ == "__main__":
    import numpy as np
    import mindspore

    # 随机输入数据
    B, C, H, W = 2, 128, 32, 32
    cost_dim = 256
    context_dim = 64
    depth_num = 1

    # 测试 DepthHead
    depth_head = DepthHead(input_dim=C, hidden_dim=64)
    x_d = mindspore.Tensor(np.random.randn(B, C, H, W), mindspore.float32)
    out = depth_head(x_d)
    print("DepthHead输入形状:", x_d.shape)
    print("DepthHead输出形状:", out.shape)

    # 测试 ConvGRU
    conv_gru = ConvGRU(hidden_dim=128, input_dim=128)
    h = mindspore.Tensor(np.random.randn(B, 128, H, W), mindspore.float32)
    x1 = mindspore.Tensor(np.random.randn(B, 64, H, W), mindspore.float32)
    x2 = mindspore.Tensor(np.random.randn(B, 64, H, W), mindspore.float32)
    out = conv_gru(h, x1,x2)
    print("ConvGRU输入形状:h:", h.shape, "x1:", x1.shape, "x2:", x2.shape)
    print("ConvGRU输出形状:", out.shape)

    # 测试 UpMaskNet
    upmask_net = UpMaskNet(hidden_dim=128, ratio=8)
    feat = mindspore.Tensor(np.random.randn(B, 128, H, W), mindspore.float32)
    mask = upmask_net(feat)
    print("UpMaskNet输入形状:", feat.shape)
    print("UpMaskNet输出形状:", mask.shape)

    # 测试 ProjectionInput
    proj_input = ProjectionInput(cost_dim=cost_dim, hidden_dim=128, context_dim=context_dim, out_chs=128, depth_num=depth_num)
    disp = mindspore.Tensor(np.random.randn(B, depth_num, H, W), mindspore.float32)
    cost = mindspore.Tensor(np.random.randn(B, cost_dim, H, W), mindspore.float32)
    context = mindspore.Tensor(np.random.randn(B, context_dim, H, W), mindspore.float32)
    out = proj_input(disp, cost, context)
    print("ProjectionInput输入形状: disp:", disp.shape, "cost:", cost.shape, "context:", context.shape)
    print("ProjectionInput输出形状:", out.shape)
    print("ProjectionInput输出形状out_chs:",proj_input.out_chs)

    # 测试 BasicUpdateBlock
    def dummy_depth_cost_func(x, iter=0):
        # 返回一个和cost_dim一样shape的Tensor
        return mindspore.Tensor(np.random.randn(B, cost_dim, H, W), mindspore.float32)
    def dummy_scale_inv_depth(x):
        # 返回一个tuple，第2项是和disp一样shape的Tensor
        min_depth = 2
        max_depth = 100
        min_disp = 1 / max_depth
        max_disp = 1 / min_depth
        scaled_disp = min_disp + (max_disp - min_disp) * disp
        scaled_disp = scaled_disp.clamp(min = 1e-4)
        depth = 1 / scaled_disp
        return scaled_disp, depth
    basic_update_block = BasicUpdateBlock(hidden_dim=128, cost_dim=cost_dim, ratio=8, context_dim=context_dim, UpMask=True, cost_num=1)
    net = mindspore.Tensor(np.random.randn(B, 128, H, W), mindspore.float32)
    inv_depth = mindspore.Tensor(np.random.randn(B, depth_num, H, W), mindspore.float32)
    context = mindspore.Tensor(np.random.randn(B, context_dim, H, W), mindspore.float32)
    print("BasicUpdateBlock输入net形状:", net.shape, "inv_depth形状:", inv_depth.shape, "context形状:", context.shape)
    out_net, mask_list, inv_depth_list = basic_update_block(
        net, 
        dummy_depth_cost_func, 
        inv_depth, 
        context, 
        seq_len=4, 
        scale_inv_depth=dummy_scale_inv_depth
    )
    print("BasicUpdateBlock输出net形状:", out_net.shape)
    print("BasicUpdateBlock输出mask_list长度:", len(mask_list))
    print("BasicUpdateBlock输出inv_depth_list长度:", len(inv_depth_list))

